from Trainer.trainer import Trainer
from tqdm import tqdm
import time

class SupervisedTrainer(Trainer):
    def train(self):
        epochs = self.train_cfg["epochs"]
        self.model.train()
        loss_save_list = []
        train_start_time = time.time()
        for epoch in range(epochs):
            running_loss = 0.0
            print(f"epoch: {epoch}")
            for i, data in enumerate(tqdm(self.train_loader), 0):
                labels = data[-1]
                labels = labels.to(self.device)
                inputs = []
                for j in range(len(data) - 1):
                    inputs.append(data[j].to(self.device))
                self.optimizer.zero_grad()
                loss = 0.0
                for input in inputs:
                    outputs = self.model(input)
                    if (isinstance(outputs, list)):
                        for i in range(self.train_cfg["num_sub_heads"]):
                            loss += self.loss(outputs[i], labels)
                    else:
                        loss += self.loss(outputs, labels)
                loss.backward()
                self.optimizer.step()
                running_loss += loss.item()
            self.writer.add_scalar(f"loss", running_loss / len(self.train_loader), epoch)
            print(f"Loss: {running_loss/ len(self.train_loader)}")
            loss_save_list.append(running_loss / len(self.train_loader))
            self.save_model(epoch)
        train_end_time = time.time()
        self.train_time = train_end_time - train_start_time
        # 保存损失记录
        self.save_loss(loss_save_list)
